from torchaudio.transforms import Spectrogram
import numpy as np
import librosa as lb

stft_configs = {
    "n_fft": 2048,
    "hop_length": 256,
    "win_length": 2048,
    "sample_rate": 16000
}

stft_transform = Spectrogram(
    n_fft=stft_configs["n_fft"],
    hop_length=stft_configs["hop_length"],
    win_length=stft_configs["win_length"],
    power=None,  # Use complex output
    normalized=True,
    center=True,
    onesided=True
    )

def stft(waveform):
    return stft_transform(waveform).abs()

def fft(waveform):
    return np.abs(np.fft.fft(waveform, axis=-1))[:waveform.shape[-1] // 2 + 1]

if __name__ == "__main__":
    # Example usage
    import soundfile as sf
    import torch
    
    wav, fs = sf.read("/nvmework1/shaonian/Datasets/AudioSet_extracted/Y__0Fp4K-2Ew.wav")
    X = stft(torch.tensor(wav).unsqueeze(0))
    print(X.shape)